# 首先导入修复脚本，解决三个主要问题：
# 1. HuggingFace transformers 的 chat_template 设置问题
# 2. BERTScore 的 meta tensor 设备迁移问题
# 3. METEOR 的 NLTK 语料库兼容问题
import sys
import os
# 将脚本目录添加到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
# 导入修复脚本
# 应用所有修复
apply_all_fixes()

import json
import logging
import argparse
import os
from typing import List, Dict, Any
from concurrent.futures import ProcessPoolExecutor, as_completed

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class MultiTurnExperimentRunner:
    """
    多轮对话实验运行器，负责批量执行多轮对话实验
    """
    
    def __init__(self, config_file_path: str):
        """
        初始化多轮对话实验运行器
        
        Args:
            config_file_path: 配置文件路径
        """
        self.config = self.load_config(config_file_path)
        
        # 确保结果目录存在
        self.results_dir = self.config.get('results_dir', '/mnt/ssd/jsj/patient/results/multi_turn')
        os.makedirs(self.results_dir, exist_ok=True)
        
        # 保存配置到结果目录
        self._save_config_to_results()
    
    def load_config(self, config_file_path: str) -> Dict[str, Any]:
        """
        加载配置文件
        
        Args:
            config_file_path: 配置文件路径
        
        Returns:
            配置字典
        """
        try:
            if not os.path.exists(config_file_path):
                logger.error(f"配置文件不存在: {config_file_path}")
                # 返回默认配置
                return self._get_default_config()
            
            with open(config_file_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            
            # 合并默认配置
            default_config = self._get_default_config()
            for key, value in default_config.items():
                if key not in config:
                    config[key] = value
            
            logger.info(f"成功加载配置文件: {config_file_path}")
            return config
        except Exception as e:
            logger.error(f"加载配置文件失败: {str(e)}")
            # 返回默认配置
            return self._get_default_config()
    
    def _get_default_config(self) -> Dict[str, Any]:
        """
        获取默认配置
        
        Returns:
            默认配置字典
        """
        return {
            'experiment_name': 'multi_turn_dialogue_experiment',
            'initial_history_turns': 2,
            'max_experiment_turns': 5,
            'context_window_size': 'full',
            'parallel_processes': 4,
            'selected_patients': [],
            'personality_combinations': [],
            'evaluation_metrics': ['bleu', 'rouge'],
            'results_dir': '/mnt/ssd/jsj/patient/results/multi_turn',
            'dataset_dir': '/mnt/ssd/jsj/patient/dataset'
        }
    
    def _save_config_to_results(self):
        """
        将配置保存到结果目录
        """
        try:
            config_save_path = os.path.join(self.results_dir, 'experiment_config.json')
            with open(config_save_path, 'w', encoding='utf-8') as f:
                json.dump(self.config, f, ensure_ascii=False, indent=2)
            logger.info(f"配置已保存到: {config_save_path}")
        except Exception as e:
            logger.error(f"保存配置失败: {str(e)}")
    
    def collect_patient_files(self) -> List[str]:
        """
        收集患者数据文件
        
        Returns:
            患者数据文件路径列表
        """
        patient_files = []
        dataset_dir = self.config.get('dataset_dir', '/mnt/ssd/jsj/patient/dataset')
        
        # 计算所需的对话轮次总数
        initial_history_turns = self.config.get('initial_history_turns', 4)
        max_experiment_turns = self.config.get('max_experiment_turns', 5)
        required_turns = initial_history_turns + max_experiment_turns
        
        logger.info(f"正在扫描数据集目录: {dataset_dir}")
        logger.info(f"筛选要求：至少包含 {required_turns} 轮对话")
        
        # 遍历dataset目录下的所有子目录
        for root, dirs, files in os.walk(dataset_dir):
            for file in files:
                # 选择基本患者数据文件（不含人格特征描述的文件）
                if file.endswith('.json') and not file.endswith('_normal.json') and not any(p in file for p in ['偏执', '啰嗦', '怀疑', '合作']):
                    file_path = os.path.join(root, file)
                    try:
                        # 读取文件内容，检查对话轮次
                        with open(file_path, 'r', encoding='utf-8') as f:
                            patient_data = json.load(f)
                        
                        # 获取对话列表
                        dialogues = patient_data.get('门诊对话', [])
                        
                        # 检查对话轮次是否满足要求
                        if len(dialogues) >= required_turns:
                            patient_files.append(file_path)
                            logger.debug(f"文件 {file_path} 满足对话轮次要求 ({len(dialogues)} >= {required_turns})")
                        else:
                            logger.debug(f"文件 {file_path} 对话轮次不足 ({len(dialogues)} < {required_turns})")
                    except Exception as e:
                        logger.error(f"读取文件 {file_path} 失败: {str(e)}")
        
        # 根据配置确定要选择的患者数量
        selected_count = self.config.get('selected_patients', 0)
        if selected_count > 0 and selected_count < len(patient_files):
            # 如果指定了患者数量，则从中随机选择指定数量的患者
            import random
            original_count = len(patient_files)
            patient_files = random.sample(patient_files, selected_count)
            logger.info(f"已从满足要求的 {original_count} 个患者中随机选择 {selected_count} 个进行实验")
        
        logger.info(f"共收集到 {len(patient_files)} 个满足对话轮次要求的患者数据文件")
        return patient_files
    
    def run_experiment_for_patient(self, patient_file_path: str) -> Dict[str, Any]:
        """
        为单个患者运行多轮对话实验
        
        Args:
            patient_file_path: 患者数据文件路径
        
        Returns:
            实验结果
        """
        try:
            # 导入MultiTurnExperimentManager
            import sys
            sys.path.append('/mnt/ssd/jsj/patient/script')
            from multi_turn_experiment_manager import MultiTurnExperimentManager
            
            # 创建多轮对话实验管理器
            multi_turn_manager = MultiTurnExperimentManager(self.config)
            
            # 运行实验
            result = multi_turn_manager.run_multi_turn_experiment(patient_file_path)
            
            return result
        except Exception as e:
            logger.error(f"为患者文件 {patient_file_path} 运行实验失败: {str(e)}")
            return {}
    
    def run_batch_experiments(self):
        """
        批量运行多轮对话实验
        """
        # 收集患者数据文件
        patient_files = self.collect_patient_files()
        if not patient_files:
            logger.error("没有找到有效的患者数据文件，无法运行实验")
            return
        
        # 根据配置确定并行进程数
        parallel_processes = min(self.config.get('parallel_processes', 4), len(patient_files))
        
        logger.info(f"开始批量运行多轮对话实验，共 {len(patient_files)} 个患者，并行进程数: {parallel_processes}")
        
        # 使用进程池并行运行实验
        results = []
        
        with ProcessPoolExecutor(max_workers=parallel_processes) as executor:
            # 提交所有任务
            futures = {executor.submit(self.run_experiment_for_patient, file_path): file_path for file_path in patient_files}
            
            # 处理完成的任务
            for future in as_completed(futures):
                patient_file = futures[future]
                try:
                    result = future.result()
                    if result:
                        results.append(result)
                        patient_id = result.get('patient_id', '未知患者')
                        logger.info(f"患者 {patient_id} 的实验已完成")
                except Exception as e:
                    logger.error(f"处理患者文件 {patient_file} 时发生异常: {str(e)}")
        
        logger.info(f"批量多轮对话实验完成，成功处理 {len(results)} 个患者")
        
        # 生成汇总报告
        self._generate_summary_report(results)
    
    def _generate_summary_report(self, results: List[Dict[str, Any]]):
        """
        生成实验汇总报告
        
        Args:
            results: 所有实验结果列表
        """
        try:
            summary = {
                'total_patients': len(results),
                'experiment_name': self.config.get('experiment_name', 'multi_turn_dialogue_experiment'),
                'config_summary': {
                    'initial_history_turns': self.config.get('initial_history_turns', 2),
                    'max_experiment_turns': self.config.get('max_experiment_turns', 5),
                    'context_window_size': self.config.get('context_window_size', 'full')
                },
                'results_summary': []
            }
            
            # 为每个患者生成结果摘要
            for result in results:
                patient_summary = {
                    'patient_id': result.get('patient_id', '未知患者'),
                    'personality': result.get('personality', {}),
                    'total_turns': result.get('total_turns', 0),
                    'metrics_summary': {}
                }
                
                # 计算每轮对话的平均指标
                if 'experiment_results' in result:
                    experiment_results = result['experiment_results']
                    if experiment_results:
                        # 初始化指标总和
                        metrics_sum = {}
                        
                        # 定义递归添加指标值的函数
                        def add_metric_value(sum_dict, metric_name, metric_value):
                            if isinstance(metric_value, dict):
                                if metric_name not in sum_dict:
                                    sum_dict[metric_name] = {}
                                for k, v in metric_value.items():
                                    add_metric_value(sum_dict[metric_name], k, v)
                            elif isinstance(metric_value, (int, float)):
                                if metric_name not in sum_dict:
                                    sum_dict[metric_name] = 0
                                sum_dict[metric_name] += metric_value
                            # 忽略其他类型的值
                        
                        for turn_result in experiment_results:
                            if 'metrics' in turn_result:
                                for metric_name, metric_value in turn_result['metrics'].items():
                                    add_metric_value(metrics_sum, metric_name, metric_value)
                        
                        # 定义递归计算平均值的函数
                        def calculate_average(metric_dict, count):
                            result_dict = {}
                            for k, v in metric_dict.items():
                                if isinstance(v, dict):
                                    result_dict[k] = calculate_average(v, count)
                                elif isinstance(v, (int, float)):
                                    result_dict[k] = v / count
                            return result_dict
                        
                        # 计算平均值
                        patient_summary['metrics_summary'] = calculate_average(metrics_sum, len(experiment_results))
                
                summary['results_summary'].append(patient_summary)
            
            # 保存汇总报告
            summary_file_path = os.path.join(self.results_dir, 'experiment_summary.json')
            with open(summary_file_path, 'w', encoding='utf-8') as f:
                json.dump(summary, f, ensure_ascii=False, indent=2)
            
            logger.info(f"实验汇总报告已保存到: {summary_file_path}")
        except Exception as e:
            logger.error(f"生成实验汇总报告失败: {str(e)}")

def main():
    """
    主函数
    """
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='运行多轮对话实验')
    parser.add_argument('--config', type=str, default='/mnt/ssd/jsj/patient/script/multi_turn_config.json',
                        help='配置文件路径，默认为/mnt/ssd/jsj/patient/script/multi_turn_config.json')
    parser.add_argument('--patient', type=str, help='单个患者数据文件路径，如果指定则只处理该患者')
    
    args = parser.parse_args()
    
    # 创建多轮对话实验运行器
    runner = MultiTurnExperimentRunner(args.config)
    
    # 运行实验
    if args.patient:
        # 只处理单个患者
        logger.info(f"开始为单个患者运行多轮对话实验: {args.patient}")
        result = runner.run_experiment_for_patient(args.patient)
        logger.info(f"单个患者多轮对话实验完成，结果: {result}")
    else:
        # 批量处理多个患者
        runner.run_batch_experiments()

if __name__ == "__main__":
    main()